
import caffe
import numpy as np
import lmdb
import os

# root_file_path = "/home/sunzy/workspace/pyhome/CaffeProject/CaffeProject/pulsar/deep_learning/temp/"

name = "subints"

root_file_path = "/home/dataology/workspace/caffe/CaffeProject/pulsar/deep_learning/temp/"
deploy = root_file_path + '%s/deploy.prototxt' % name

def test(model_id, niter):
    """
    test the model and return the P,R
    :param model_id:the id of the model(hyper parameter)
    :param niter: number of iteration
    :return: precision and recall of pulsars
    """
    caffe_model = root_file_path + '%s/%d_custom_net_iter_%d.caffemodel' % (name, model_id, niter)
    net = caffe.Net(deploy, caffe_model, caffe.TEST)
    # lmdb_path = '/home/sunzy/workspace/data/MedlatTrainingData/lmdb'
    lmdb_path = '/usr/data/MedlatTrainingData/lmdb'
    img_lmdb = lmdb.open(os.path.join(lmdb_path, "%s_test_lmdb" % name))
    txn = img_lmdb.begin()
    cursor = txn.cursor()
    datum = caffe.proto.caffe_pb2.Datum()
    num = 0
    real_p = 0
    tp = 0
    pre_p = 0
    for (idx, (key, value)) in enumerate(cursor):
        datum.ParseFromString(value)
        flat_x = np.fromstring(datum.data, dtype=np.float64)
        # print(key)
        x = flat_x.reshape(datum.channels, datum.height, datum.width)
        y = datum.label
        num += 1
        # print(x*255)
        # if num == 1:
        #     break

        net.blobs['data'].data[...] = x
        out = net.forward()

        prob = out['prob']
        order = prob.argmax()
        print(num, y, order)

        # print(num, order, prob)
        if y == 1:
            real_p += 1
        if order == 1:
            pre_p += 1
        if order == 1 and y == 1:
            tp += 1

    print(num, real_p, tp, pre_p)
    # return precision and recall
    return tp/(pre_p+1e-10), tp/(real_p+1e-10)



def main():
    p, r = test(0, 2000)
    print(p, r)

if __name__ == "__main__":
    main()